# 导入tensorflow库
import torch
# 导入matplotlib.pyplot库，用于绘图
import matplotlib.pyplot as plt
# 导入numpy库，用于数组操作
import numpy as np

# 设置绘图风格
plt.style.use("seaborn")

# 定义输入数据x和目标数据y
x = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0])
y = torch.tensor([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0])

# 声明线性模型 y = wx+b 中的 w 和 b
w = torch.tensor(-3.0, requires_grad=True)
b = torch.tensor(10.0, requires_grad=True)
# 设置学习率为0.09
lr = 0.09

# 初始化用于记录梯度下降过程中的w,b,以及对应的损失的列表
w_list = []
b_list = []
loss_list = []
# 初始化循环次数
epoch = 0
while True:
    epoch = epoch + 1  # 记录循环次数
    # 计算当前 w 和 b 对应的损失，并记录
    y_guess = w * x + b
    loss = torch.mean((y_guess - y) ** 2)
    # 将损失值添加到损失列表中
    loss.backward()

    loss_list.append(loss.item())
    # 计算当前 w 和 b 的梯度
    # 将当前的 w 和 b 添加到对应的列表中

    w_list.append(w.item())
    b_list.append(b.item())
    # 调节 w 和 b 实现梯度下降的过程
    # 注意：必须把结果转化为 tensor，否则运算会出错
    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad
        w.grad.zero_()
        b.grad.zero_()

    # 如果 loss 足够小，则退出循环
    if loss < 1e-5:
        break
y_guess = y_guess.detach().numpy()

# 将循环次数转换为range类型
x = range(epoch)
# 打印迭代次数和最终的 w 和 b 值
print(f"一共尝试了{epoch}次迭代周期(epoch)")
print(f"最终得到: w={w.item():.10f} b={b.item():.10f}")
print(f"四舍五入: w={w.item():.0f} b={b.item():.0f}")

# 创建一个1行4列的子图
fg, ax = plt.subplots(1, 4, figsize=(12, 3))

# 在第一个子图中绘制 w 的变化过程
ax[0].plot(x, w_list)
ax[0].set_title("w → 2", size=16)

# 在第二个子图中绘制 b 的变化过程
ax[1].plot(x, b_list)
ax[1].set_title("b → -1", size=16)

# 在第三个子图中绘制损失的变化过程
ax[2].plot(x, loss_list)
ax[2].set_title("loss", size=16)

# 在第四个子图中绘制真实值和预测值的对比
x = np.array([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0])
ax[3].scatter(x, y, color='black')

ax[3].scatter(x, y_guess, color='red', s=20)
ax[3].vlines(x, y, y_guess, linestyles='dotted', color='red')
ax[3].plot(x, y, '--', color='blue')
ax[3].plot(x, y_guess, 'r', linewidth=0.9)
ax[3].set_title(f"y={w.item():.0f}x+{b.item():.0f}", size=20)
ax[3].legend(['real', 'guess'], loc='upper left')

# 调整子图的布局并显示
plt.tight_layout()
plt.show()

